
from visual import *
import ode
from QuadMotor import *
from QuadPropeller import *
from PIDController import *

from decimal import *

def rotmatrix_to_yawpitchroll(rot_matrix):
    yaw = 0.0
    roll = 0.0
    pitch = asin(-rot_matrix[9])

    threshold = 0.001

    test = cos(pitch)

    if test > threshold:
        roll = atan2(rot_matrix[1], rot_matrix[5])
        yaw = atan2(rot_matrix[8], rot_matrix[10])

    else:
        roll = atan2(-rot_matrix[4], rot_matrix[0])
        yaw = 0.0

    #getcontext().prec = 12

    #EIGHTPLACES = Decimal(10) ** -8

    ret_yaw = Decimal(str(yaw))
    ret_pitch = Decimal(str(pitch))
    ret_roll = Decimal(str(roll))

    return ret_yaw,ret_pitch,ret_roll

def QuaternoinToYawPitchRoll(odequat):
    epsilon = 0.009765625
    threshold = 0.5 - epsilon
    qx = odequat[1]
    qy = odequat[2]
    qz = odequat[3]
    qw = odequat[0]
    yaw = 0
    pitch = 0
    roll = 0

    xy = qx * qy
    zw = qz * qw

    test = xy + zw
    if test < - threshold or test > threshold:
        sign = 0
        if test < 0:
            sign = -1
        if test > 0 :
            sign = 1
        yaw = sign * 2 * atan2(qx, qw)
        pitch = sign * (3.14159 / 2.0)
        roll = 0

    else:
        xx  = qx * qx
        xz = qx * qz
        xw = qx * qw

        yy = qy * qy
        yw = qy * qw
        yz = qy * qz

        zz = qz * qz

        yaw = atan2( 2 * yw - 2 * xz , 1 - 2 * yy - 2 * zz)
        pitch = atan2( 2 * xw - 2 * yz, 1 - 2 * xx - 2 * zz)
        roll = asin(2 * test)

    return yaw, pitch, roll

class QuadRotor(object):
    def __init__(self, world, space):
        self.quad_frame = frame()
        self.vis_body = box(frame = self.quad_frame, pos = (0,0,0), length = 5, width = 5, height = 2)
        self.vis_body.color = color.red
        self.body = ode.Body(world)
        body_mass = ode.Mass()
        body_mass.setBox(20,5,2,5)
        body_mass.mass = 1.0

        self.body.setPosition((0,20,0))
        self.body_geom = ode.GeomBox(space, lengths=(5,2,5))
        self.body_geom.setBody(self.body)

        self.tframe_mass, self.tframe_geom = self.build_frame(space, (0,0,5), (1, 1, 10), color.red)
        self.bframe_mass, self.bframe_geom = self.build_frame(space, (0,0,-5), (1,1,10), color.red)
        self.rframe_mass, self.rframe_geom = self.build_frame(space, (5,0,0), (10,1,1), color.red)
        self.lframe_mass, self.lframe_geom = self.build_frame(space, (-5,0,0), (10,1,1), color.red)

        body_mass.add(self.tframe_mass)
        body_mass.add(self.bframe_mass)
        body_mass.add(self.rframe_mass)
        body_mass.add(self.lframe_mass)



        self.motors = []
        #1 is counterclockwise and -1 is clockwise
        #front motor
        self.motors.append(QuadMotor(self.quad_frame, self.body, space, (0,-0.5,10.9), 1,1, color.blue))
        #back motor
        self.motors.append(QuadMotor(self.quad_frame, self.body, space, (0,-0.5,-10.9), 1,1, color.blue))
        #right motor
        self.motors.append(QuadMotor(self.quad_frame, self.body, space, (10.9,-0.5,0), 1,-1, color.blue))
        #left motor
        self.motors.append(QuadMotor(self.quad_frame, self.body, space, (-10.9,-0.5,0), 1,-1, color.blue))


        for motor in self.motors:
            motor.attach_propeller(world,QuadPropeller(world,space,motor.get_pos() + vector(0,22,0),(0.1,0.1,5),color.yellow))
            body_mass.add(motor.get_mass())

        self.body.setMass(body_mass)

        self.init_rot = self.body.getRotation()
        self.init_pos = self.body.getPosition()

        self.thrust = 0.0
        self.yaw = 0.0

        yaw, pitch, roll = QuaternoinToYawPitchRoll(self.body.getQuaternion())
        #yaw, pitch, roll = rotmatrix_to_yawpitchroll(self.body.getRotation())

        self.thrust_pid = PIDController()
        self.roll_pid = PIDController()
        self.roll_pid.ziegler_nicholas_tune(Decimal('0.5'), Decimal('1') , self.roll_pid.TEST_PID)
        self.roll_pid.set_actual(roll)
        self.pitch_pid = PIDController()
        self.pitch_pid.ziegler_nicholas_tune(Decimal('0.5'), Decimal('1') , self.pitch_pid.TEST_PID)
        self.pitch_pid.set_actual(pitch)
        self.yaw_pid = PIDController()
        self.yaw_pid.ziegler_nicholas_tune(Decimal('0.5'), Decimal('1') , self.yaw_pid.TEST_PID)
        self.yaw_pid.set_actual(yaw)

        self.thrust_pid.set_target(self.init_pos[0])
        self.thrust_pid.set_actual(self.init_pos[0])

        self.thrust_pid.compute_prev_error()
        self.roll_pid.compute_prev_error()
        self.yaw_pid.compute_prev_error()
        self.pitch_pid.compute_prev_error()

    def get_motor_dir(self, motor_num):
        return self.motors[motor_num].get_dir()

    def set_motor_dir(self, motor_num, motor_dir):
        self.motors[motor_num].set_dir(motor_dir)

    def set_thrust(self, thrust):
        self.thrust = thrust
        self.thrust_pid.set_target(self.body.getPosition()[1]+self.thrust)

    def set_yaw(self, yaw):
        self.yaw = yaw
        self.yaw_pid.set_target(yaw)

    def reset(self):
        self.body.setRotation(self.init_rot)
        self.body.setPosition(self.init_pos)
        self.body.setForce((0,0,0))
        self.body.setTorque((0,0,0))
        self.body.setLinearVel((0,0,0))
        self.body.setAngularVel((0,0,0))
        self.thrust = 0
        self.yaw = 0
        self.thrust_pid.reset()
        self.yaw_pid.reset()
        self.pitch_pid.reset()
        self.roll_pid.reset()

        for motor in self.motors:
            motor.set_speed(self.thrust)



    def build_frame(self, space, frame_pos, size, color):
        vis_frame = box(frame = self.quad_frame, pos = frame_pos, length = size[0], width = size[2], height = size[1])
        vis_frame.color = color;
        frame_mass = ode.Mass()
        frame_mass.setBox(20,size[0],size[1],size[2])
        frame_mass.mass = 1.0
        frame_mass.translate(frame_pos)

        frame_geom = ode.GeomBox(lengths = size)
        frame_geom.setPosition(frame_pos)

        frame_transform = ode.GeomTransform(space)
        frame_transform.setGeom(frame_geom)

        frame_transform.setBody(self.body)

        return frame_mass, frame_transform

    def update(self):
        self.quad_frame.pos = self.body.getPosition()
        rot = self.body.getRotation()
        self.quad_frame.axis=rot[0], rot[3], rot[6]
        self.quad_frame.up=rot[1], rot[4], rot[7]

        front_motor = self.motors[0]
        back_motor = self.motors[1]
        right_motor = self.motors[2]
        left_motor = self.motors[3]

        yaw, pitch, roll = QuaternoinToYawPitchRoll(self.body.getQuaternion())


        self.thrust_pid.update(self.quad_frame.pos[1])
        self.yaw_pid.update(yaw)
        self.pitch_pid.update(pitch)
        self.roll_pid.update(roll)

        new_thrust = self.thrust_pid.get_output()
        new_pitch = max(min(self.pitch_pid.get_output(), 100), -100)
        new_yaw = max(min(self.yaw_pid.get_output(), 100) , -100)
        new_roll = max(min(self.roll_pid.get_output(), 100), -100)

        print str(new_yaw)

        front_motor.set_speed(new_thrust + new_pitch + new_yaw)
        left_motor.set_speed(new_thrust + new_roll - new_yaw)
        right_motor.set_speed(new_thrust - new_roll - new_yaw)
        back_motor.set_speed(new_thrust - new_pitch + new_yaw)

        for motor in self.motors:
            motor.update()



